#!/usr/bin/env python

#import ssummo.ssummolib as slib

from ssummo.cmd_options import Options
from ssummo import colours
from ssummo import CONFIG
from ssummo.ssummolib import load_index
from ssummo.taxonomy import TaxDB, fetch_rank
from ssummo.traverse import find_node_name
from ssummo import phyloxml

import os
import re
import sys

class LocalOptions( Options ):
    def __init__(self, args=None):
        self.options = {
            '-out'       : None,
            '-in'        : '',
            '-itol-type' : 'multivaluebar',
            '-name-col'  : 0,
            '-data-col'  : ['all'],
            '--headers'  : True,
            '-headers'   : [],
            '-delimiter' : '\t',
        }
        self.singleargs = [
            '-in',
            '-out',
            '-graph-type',
            '-data-col',
            '-name-col',
        ]
        self.multiargs = ['-headers', '-data-col']
        self.help_text = {
            '-itol-type' : 'IToL figure type. Currently only multivaluebar supported',
            '-in'        : '',
            '-out'       : 'Output file for prefix for phyloxml file and IToL data file',
            '-name-col'  : 'Column index which contains the name (first col is 0).',
            '-data-col'  : 'Space separated list of column numbers which contain data (first is 0).',
            '--headers'  : 'If the file contains headers, use them in the field descriptions [True]',
            '-headers'   : 'If the file does not contain headers, supply header info here',
            '-delimiter' : 'Data delimiter',
        }
        self.post_checks += [self.integers, self.reqs]
        if args is not None:
            self.parse_args(args)

    def integers(self):
        self['-name-col'] = int(self['-name-col'])
        if 'all' not in self['-data-col']:
            for col, entry in enumerate(self['-data-col']):
                self['-data-col'][col] = int(entry)

    def reqs(self):
        if self['-out'] is None:
            raise IOError( 'Need to define -out!\n' )

def separate( options ):
    file_name = options['-in']
    with file( file_name, 'r' ) as in_file:
        for line in in_file:
            row = line.rstrip().split( options['-delimiter'] )
            yield row

def join_results( results, taxpath, species ):
    path = taxpath.split( os.path.sep )
    if len( path ) >= 2:
        path = path[:-1]
    node = results
    for tax in path:
        if tax in node:
            node = node[tax]
        else:
            node.update( { tax : {} } )
            node = node[tax]
    node.update( { species : {} } )
    return results

def main(options):
    TDB = TaxDB()
    tdict = load_index()
    tdict_node = tdict
    splitter = separate(options)
    spaceSub = re.compile(r'\s+')
    text_formatter = re.compile( '["\',/]' )
    with file(options['-out'] + '.txt', 'w' ) as out_file:
        headers = splitter.next()
        if options['-data-col'] == ['all']:
            options['-data-col'] = range(0, options['-name-col']) + \
                                   range(1+options['-name-col'], len(headers))
        get_data = lambda data: [data[index] for index in options['-data-col']]
        print( get_data(headers) )
        out_file.write('LABELS,{0}\n'.format(','.join(get_data(headers))))
        cols = colours.generate_HEX_colours( len(options['-data-col']) )
        out_file.write( 'COLORS,{0}\n'.format( ','.join(cols) ) )
        results = {}
        for line in splitter:
            species = text_formatter.sub('', line[options['-name-col']])
            # MAKE DICTIONARY
            print('Searching for {0}'.format(species)),
            (rank, name) = fetch_rank(TDB, 'Prokaryotes', species)
            print('. Found {0} at rank: {1}'.format(name, rank))
            if rank is not False:
                taxpath = find_node_name(name, tdict_node)
                results = join_results(results, taxpath, species)

                data = get_data(line)
                out_file.write('{0},{1}\n'.format(
                               spaceSub.sub('_', species), ','.join(data)))
            #data = []
            #for col in options['-data-col']:
            #    data.append( line[col] )
    with file(options['-out'] + '.xml', 'w') as xmlfile:
        print('Saving phyloxml file to {0}'.format(xmlfile.name))
        phyloxml.write_xml(CONFIG.top, xmlfile, results, tdict, )

if __name__ == '__main__':
    args = sys.argv[1:]
    options = LocalOptions(args)
    main( options )
